import itk
import vtk
import time
import numpy as np

import connect2


class test():
    #
    # Test the built in ITK->VTK function
    #
    def __init__(self):
        self.vtk2itk=itk.VTKImageToImageFilter.IUS2.New()
        self.itk2vtk=itk.ImageToVTKImageFilter.IUS2.New()

        self.my_vtk2itk = connect2.vtk2itk()
        self.my_itk2vtk = connect2.itk2vtk()

        self.fixedImageName = 'C:/Users/bryan/bryan-code/2D3D/vert1/fluoro/ushortim080-LAT.mhd'
         
        self.itk_reader = itk.ImageFileReader.IUS2.New()
        self.itk_reader.SetFileName(self.fixedImageName)
        self.itk_reader.Update()

        self.vtk_reader = vtk.vtkMetaImageReader()
        self.vtk_reader.SetFileName(self.fixedImageName)
        self.vtk_reader.Update()


    def connect_builtin(self):
        start = time.clock()

        self.vtk2itk.SetInput(self.vtk_reader.GetOutput())
        itk_img = self.vtk2itk.GetOutput()
        itk_img.Update()
        
        self.itk2vtk.SetInput(self.itk_reader.GetOutput())
        vtk_img = self.itk2vtk.GetOutput()
        vtk_img.Update()

        end = time.clock()

        return end-start

    def connect_my(self):
        start = time.clock()

        self.my_vtk2itk.SetInput(self.vtk_reader.GetOutput())
        itk_img = self.my_vtk2itk.GetOutput()
        itk_img.Update()
        
        self.my_itk2vtk.SetInput(self.itk_reader.GetOutput())
        vtk_img = self.my_itk2vtk.GetOutput()
        vtk_img.Update()

        end = time.clock()

        return end-start


    def run_test(self):
        time1 = self.connect_builtin()
        time2 = self.connect_my()

        print "Built in function: %f, My function: %f" % (time1,time2)
        return time1,time2


if __name__ == "__main__":

    this = test()
    a=[]
    b=[]
    for i in range(100):
        time1,time2=this.run_test()
        a.append(time1)
        b.append(time2)
        
    a=np.asarray(a)
    b=np.asarray(b)
    print "===AVERAGE==="
    print "Built in function: %f, My function: %f" % (a.mean(),b.mean())